package com.ycy.lucene.plugin;

import org.apache.commons.lang3.StringUtils;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.FieldType;
import org.apache.lucene.document.IntPoint;
import org.apache.lucene.document.StringField;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.IndexWriterConfig.OpenMode;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.Terms;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.search.BooleanClause;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.BoostQuery;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TermQuery;
import org.apache.lucene.search.spell.Dictionary;
import org.apache.lucene.search.spell.LevenshteinDistance;
import org.apache.lucene.search.spell.StringDistance;
import org.apache.lucene.search.spell.SuggestMode;
import org.apache.lucene.search.spell.SuggestWord;
import org.apache.lucene.search.spell.SuggestWordQueue;
import org.apache.lucene.store.AlreadyClosedException;
import org.apache.lucene.store.Directory;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.BytesRefIterator;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;

/**
 * 基于索引格式纠错
 * 注意：文件格式区别于lucene原依赖格式（word1/n word2/n ...)
 * 全站纠错文件格式：segs1 word1/n segs2 word2/n
 * segs是针对word的各种拆分形式，例如拼音、英文名，偏旁自首
 * @author c_jia
 */
public class CustomSpellChecker implements java.io.Closeable {
  private static final float DEFAULT_ACCURACY = 0.8f;
  private static final String F_WORD = "word";
  private static final String F_LEN = "len";

  Directory spellIndex;
  private float bStart = 2.0f;
  private float bEnd = 1.0f;

  private IndexSearcher searcher;
  private final Object searcherLock = new Object();
  private final Object modifyCurrentIndexLock = new Object();

  private volatile boolean closed = false;
  private float accuracy = DEFAULT_ACCURACY;

  private StringDistance sd;
  private Comparator<SuggestWord> comparator;

  public CustomSpellChecker(Directory spellIndex, StringDistance sd) throws IOException {
    this(spellIndex, sd, SuggestWordQueue.DEFAULT_COMPARATOR);
  }

  public CustomSpellChecker(Directory spellIndex) throws IOException {
    this(spellIndex, new LevenshteinDistance());
  }

  public CustomSpellChecker(Directory spellIndex, StringDistance sd, Comparator<SuggestWord> comparator) throws IOException {
    setSpellIndex(spellIndex);
    setStringDistance(sd);
    this.comparator = comparator;
  }

  public void setSpellIndex(Directory spellIndexDir) throws IOException {
    // this could be the same directory as the current spellIndex
    // modifications to the directory should be synchronized
    synchronized (modifyCurrentIndexLock) {
      ensureOpen();
      if (!DirectoryReader.indexExists(spellIndexDir)) {
          IndexWriter writer = new IndexWriter(spellIndexDir,
            new IndexWriterConfig(null));
          writer.close();
      }
      swapSearcher(spellIndexDir);
    }
  }

  public void setComparator(Comparator<SuggestWord> comparator) {
    this.comparator = comparator;
  }

  public Comparator<SuggestWord> getComparator() {
    return comparator;
  }

  public void setStringDistance(StringDistance sd) {
    this.sd = sd;
  }

  public StringDistance getStringDistance() {
    return sd;
  }

  public void setAccuracy(float acc) {
    this.accuracy = acc;
  }

  public float getAccuracy() {
    return accuracy;
  }

  public String[] suggestSimilar(String word, int numSug) throws IOException {
    return this.suggestSimilar(word, numSug, null, null, SuggestMode.SUGGEST_WHEN_NOT_IN_INDEX);
  }

  public String[] suggestSimilar(String word, int numSug, float accuracy) throws IOException {
    return this.suggestSimilar(word, numSug, null, null, SuggestMode.SUGGEST_WHEN_NOT_IN_INDEX, accuracy);
  }

  public String[] suggestSimilar(String word, int numSug, IndexReader ir,
                                 String field, SuggestMode suggestMode) throws IOException {
    return suggestSimilar(word, numSug, ir, field, suggestMode, this.accuracy);
  }

  private boolean containsWitSpace(String s1, String s2) {
    if (StringUtils.isEmpty(s1) || StringUtils.isEmpty(s2)) {
      return false;
    }
    return StringUtils.contains(s1.replace("\\s", ""), s2.replace("\\s", ""));
  }

  @SuppressWarnings("deprecation")
  public String[] suggestSimilar(String word, int numSug, IndexReader ir,
                                 String field, SuggestMode suggestMode, float accuracy) throws IOException {
    return suggestSimilar(word, numSug, ir, field, suggestMode, accuracy, accuracy);
  }

  public String[] suggestSimilar(Param param) throws IOException {
    return suggestSimilar(param.word, param.numSug, param.ir, param.field, param.suggestMode, param.accuracy, param.threshold);
  }

  @SuppressWarnings("deprecation")
  public String[] suggestSimilar(String word, int numSug, IndexReader ir,
                                 String field, SuggestMode suggestMode, float accuracy, float threshold) throws IOException {
    // obtainSearcher calls ensureOpen
    final IndexSearcher indexSearcher = obtainSearcher();
    try {
      if (ir == null || field == null) {
        suggestMode = SuggestMode.SUGGEST_ALWAYS;
      }
      if (suggestMode == SuggestMode.SUGGEST_ALWAYS) {
        ir = null;
        field = null;
      }

      final int lengthWord = word.length();

      final int freq = (ir != null && field != null) ? ir.docFreq(new Term(field, word)) : 0;
      final int goalFreq = suggestMode==SuggestMode.SUGGEST_MORE_POPULAR ? freq : 0;
      // if the word exists in the real index and we don't care for word frequency, return the word itself
      if (suggestMode==SuggestMode.SUGGEST_WHEN_NOT_IN_INDEX && freq > 0) {
        return new String[] { word };
      }

      BooleanQuery.Builder query = new BooleanQuery.Builder();
      String[] grams;
      String key;

      for (int ng = getMin(lengthWord); ng <= getMax(lengthWord); ng++) {

        key = "gram" + ng; // form key

        grams = formGrams(word, ng); // form word into ngrams (allow dups too)

        if (grams.length == 0) {
          continue; // hmm
        }

        if (bStart > 0) { // should we boost prefixes?
          add(query, "start" + ng, grams[0], bStart); // matches start of word

        }
        if (bEnd > 0) { // should we boost suffixes
          add(query, "end" + ng, grams[grams.length - 1], bEnd); // matches end of word

        }
        for (int i = 0; i < grams.length; i++) {
          add(query, key, grams[i]);
        }
      }

      int maxHits = 10 * numSug;
      @SuppressWarnings("deprecation")

      int lenRange = word.length() > 6 ? 3 : 2;

      Query lenQuery = IntPoint.newRangeQuery(F_LEN, word.length() - lenRange, word.length() + lenRange);
      query.add(lenQuery, BooleanClause.Occur.MUST);
      ScoreDoc[] hits = indexSearcher.search(query.build(), maxHits).scoreDocs;

      SuggestWordQueue sugQueue = new SuggestWordQueue(numSug, comparator);

      int stop = Math.min(hits.length, maxHits);
      SuggestWord sugWord = new SuggestWord();
      for (int i = 0; i < stop; i++) {
        sugWord.string = indexSearcher.doc(hits[i].doc).get(F_WORD); // get orig word
        sugWord.score = sd.getDistance(word, sugWord.string);

        if ((sugWord.score < accuracy) && !(hits[i].score > 10) && (!containsWitSpace(sugWord.string, word))) {
          continue;
        }
        if (ir != null && field != null) { // use the user index
          sugWord.freq = ir.docFreq(new Term(field, sugWord.string)); // freq in the index
          if ((suggestMode==SuggestMode.SUGGEST_MORE_POPULAR && goalFreq > sugWord.freq) || sugWord.freq < 1) {
            continue;
          }
        }
        sugQueue.insertWithOverflow(sugWord);
        if (sugQueue.size() == numSug) {
          accuracy = sugQueue.top().score;
        }
        sugWord = new SuggestWord();
      }

      // convert to array string
      String[] list = new String[sugQueue.size()];
      for (int i = sugQueue.size() - 1; i >= 0; i--) {
        SuggestWord pop = sugQueue.pop();
        if(pop.score < threshold) {
          continue;
        }
        list[i] = pop.string;
      }
      return list;
    } finally {
      releaseSearcher(indexSearcher);
    }
  }
  /**
   * Add a clause to a boolean query.
   */
  private static void add(BooleanQuery.Builder q, String name, String value, float boost) {
    Query tq = new TermQuery(new Term(name, value));
    q.add(new BooleanClause(new BoostQuery(tq, boost), BooleanClause.Occur.SHOULD));
  }

  /**
   * Add a clause to a boolean query.
   */
  private static void add(BooleanQuery.Builder q, String name, String value) {
    q.add(new BooleanClause(new TermQuery(new Term(name, value)), BooleanClause.Occur.SHOULD));
  }

  /**
   * Form all ngrams for a given word.
   * @param text the word to parse
   * @param ng the ngram length e.g. 3
   * @return an array of all ngrams in the word and note that duplicates are not removed
   */
  private static String[] formGrams(String text, int ng) {
    int len = text.length();
    String[] res = new String[len - ng + 1];
    for (int i = 0; i < len - ng + 1; i++) {
      res[i] = text.substring(i, i + ng);
    }
    return res;
  }

  /**
   * Removes all terms from the spell check index.
   * @throws IOException If there is a low-level I/O error.
   * @throws AlreadyClosedException if the Spellchecker is already closed
   */
  public void clearIndex() throws IOException {
    synchronized (modifyCurrentIndexLock) {
      ensureOpen();
      final Directory dir = this.spellIndex;
      final IndexWriter writer = new IndexWriter(dir, new IndexWriterConfig(null)
          .setOpenMode(OpenMode.CREATE));
      writer.close();
      swapSearcher(dir);
    }
  }

  /**
   * Check whether the word exists in the index.
   * @param word word to check
   * @throws IOException If there is a low-level I/O error.
   * @throws AlreadyClosedException if the Spellchecker is already closed
   * @return true if the word exists in the index
   */
  public boolean exist(String word) throws IOException {
    // obtainSearcher calls ensureOpen
    final IndexSearcher indexSearcher = obtainSearcher();
    try{
      // TODO: we should use ReaderUtil+seekExact, we dont care about the docFreq
      // this is just an existence check
      return indexSearcher.getIndexReader().docFreq(new Term(F_WORD, word)) > 0;
    } finally {
      releaseSearcher(indexSearcher);
    }
  }

  /**
   * Indexes the data from the given {@link Dictionary}.
   * @param dict Dictionary to index
   * @param config {@link IndexWriterConfig} to use
   * @param fullMerge whether or not the spellcheck index should be fully merged
   * @throws AlreadyClosedException if the Spellchecker is already closed
   * @throws IOException If there is a low-level I/O error.
   */
  public final void indexDictionary(Dictionary dict, IndexWriterConfig config, boolean fullMerge) throws IOException {
    synchronized (modifyCurrentIndexLock) {
      ensureOpen();
      final Directory dir = this.spellIndex;
      final IndexWriter writer = new IndexWriter(dir, config);
      IndexSearcher indexSearcher = obtainSearcher();
      final List<TermsEnum> termsEnums = new ArrayList<>();

      final IndexReader reader = searcher.getIndexReader();
      if (reader.maxDoc() > 0) {
        for (final LeafReaderContext ctx : reader.leaves()) {
          Terms terms = ctx.reader().terms(F_WORD);
          if (terms != null)
            termsEnums.add(terms.iterator());
        }
      }

      boolean isEmpty = termsEnums.isEmpty();

      try {
        BytesRefIterator iter = dict.getEntryIterator();
        BytesRef currentTerm;

        terms: while ((currentTerm = iter.next()) != null) {

          String line = currentTerm.utf8ToString();
          int len = line.length();
          if (len < 3) {
            continue; // too short we bail but "too long" is fine...
          }

          if (!isEmpty) {
            for (TermsEnum te : termsEnums) {
              if (te.seekExact(currentTerm)) {
                continue terms;
              }
            }
          }

          // ok index the word
          Document doc = createDocument(line, getMin(len), getMax(len));
          writer.addDocument(doc);
        }
      } finally {
        releaseSearcher(indexSearcher);
      }
      if (fullMerge) {
        writer.forceMerge(1);
      }
      // close writer
      writer.close();
      // TODO: this isn't that great, maybe in the future SpellChecker should take
      // IWC in its ctor / keep its writer open?

      // also re-open the spell index to see our own changes when the next suggestion
      // is fetched:
      swapSearcher(dir);
    }
  }

  private static int getMin(int l) {
    if (l > 5) {
      return 3;
    }
    if (l == 5) {
      return 2;
    }
    return 1;
  }

  private static int getMax(int l) {
    if (l > 5) {
      return 4;
    }
    if (l == 5) {
      return 3;
    }
    return 2;
  }

  private static Document createDocument(String text, int ng1, int ng2) {
    Document doc = new Document();
    doc.add(new StringField(F_WORD, text, Field.Store.YES));
    doc.add(new IntPoint(F_LEN, text.length()));
    addGram(text, doc, ng1, ng2);
    return doc;
  }

  private static void addGram(String text, Document doc, int ng1, int ng2) {
    int len = text.length();
    for (int ng = ng1; ng <= ng2; ng++) {
      String key = "gram" + ng;
      String end = null;
      for (int i = 0; i < len - ng + 1; i++) {
        String gram = text.substring(i, i + ng);
        FieldType ft = new FieldType(StringField.TYPE_NOT_STORED);
        ft.setIndexOptions(IndexOptions.DOCS_AND_FREQS);
        Field ngramField = new Field(key, gram, ft);
        // spellchecker does not use positional queries, but we want freqs
        // for scoring these multivalued n-gram fields.
        doc.add(ngramField);
        if (i == 0) {
          // only one term possible in the startXXField, TF/pos and norms aren't needed.
          Field startField = new StringField("start" + ng, gram, Field.Store.NO);
          doc.add(startField);
        }
        end = gram;
      }
      if (end != null) { // may not be present if len==ng1
        // only one term possible in the endXXField, TF/pos and norms aren't needed.
        Field endField = new StringField("end" + ng, end, Field.Store.NO);
        doc.add(endField);
      }
    }
  }

  private IndexSearcher obtainSearcher() {
    synchronized (searcherLock) {
      ensureOpen();
      searcher.getIndexReader().incRef();
      return searcher;
    }
  }

  private void releaseSearcher(final IndexSearcher aSearcher) throws IOException {
      // don't check if open - always decRef
      // don't decrement the private searcher - could have been swapped
      aSearcher.getIndexReader().decRef();
  }

  private void ensureOpen() {
    if (closed) {
      throw new AlreadyClosedException("Spellchecker has been closed");
    }
  }

  /**
   * Close the IndexSearcher used by this SpellChecker
   * @throws IOException if the close operation causes an {@link IOException}
   * @throws AlreadyClosedException if the {@link  } is already closed
   */
  @Override
  public void close() throws IOException {
    synchronized (searcherLock) {
      ensureOpen();
      closed = true;
      if (searcher != null) {
        searcher.getIndexReader().close();
      }
      searcher = null;
    }
  }

  private void swapSearcher(final Directory dir) throws IOException {
    /*
     * opening a searcher is possibly very expensive.
     * We rather close it again if the Spellchecker was closed during
     * this operation than block access to the current searcher while opening.
     */
    final IndexSearcher indexSearcher = createSearcher(dir);
    synchronized (searcherLock) {
      if(closed){
        indexSearcher.getIndexReader().close();
        throw new AlreadyClosedException("Spellchecker has been closed");
      }
      if (searcher != null) {
        searcher.getIndexReader().close();
      }
      // set the spellindex in the sync block - ensure consistency.
      searcher = indexSearcher;
      this.spellIndex = dir;
    }
  }

  /**
   * Creates a new read-only IndexSearcher
   * @param dir the directory used to open the searcher
   * @return a new read-only IndexSearcher
   * @throws IOException f there is a low-level IO error
   */
  // for testing purposes
  IndexSearcher createSearcher(final Directory dir) throws IOException {
    return new IndexSearcher(DirectoryReader.open(dir));
  }

  /**
   * Returns <code>true</code> if and only if the {@link  } is
   * closed, otherwise <code>false</code>.
   *
   * @return <code>true</code> if and only if the {@link  } is
   *         closed, otherwise <code>false</code>.
   */
  boolean isClosed(){
    return closed;
  }

  public static class Param {
    String word;
    int numSug;
    IndexReader ir;
    String field;
    SuggestMode suggestMode;
    float accuracy;
    float threshold;

    public static final class ParamBuilder {
      String word;
      int numSug;
      IndexReader ir = null;
      String field = null;
      SuggestMode suggestMode = SuggestMode.SUGGEST_WHEN_NOT_IN_INDEX;
      float accuracy = DEFAULT_ACCURACY;
      float threshold = 0;

      private ParamBuilder() {
      }

      public static ParamBuilder aParam() {
        return new ParamBuilder();
      }

      public ParamBuilder word(String word) {
        this.word = word;
        return this;
      }

      public ParamBuilder numSug(int numSug) {
        this.numSug = numSug;
        return this;
      }

      public ParamBuilder ir(IndexReader ir) {
        this.ir = ir;
        return this;
      }

      public ParamBuilder field(String field) {
        this.field = field;
        return this;
      }

      public ParamBuilder suggestMode(SuggestMode suggestMode) {
        this.suggestMode = suggestMode;
        return this;
      }

      public ParamBuilder accuracy(float accuracy) {
        this.accuracy = accuracy;
        return this;
      }

      public ParamBuilder threshold(float threshold) {
        this.threshold = threshold;
        return this;
      }

      public Param build() {
        Param param = new Param();
        param.threshold = this.threshold;
        param.field = this.field;
        param.word = this.word;
        param.numSug = this.numSug;
        param.accuracy = this.accuracy;
        param.ir = this.ir;
        param.suggestMode = this.suggestMode;
        return param;
      }
    }
  }
}
